# data_manager.py
import os

from sqlalchemy.dialects.postgresql import Any

from sheng.stock.basic_data_v1.database.stock_database import StockDatabase
from sheng.stock.basic_data_v1.config.database_config import DatabaseConfig
from sheng.stock.basic_data_v1.data_sources.akshare_source import AkShareDataSource
from sheng.stock.basic_data_v1.data_sources.yahoo_source import YahooDataSource
from sheng.stock.basic_data_v1.data_sources.baostock_source import BaostockDataSource
import pandas as pd
import logging
from typing import List, Dict
import schedule
import time
from datetime import datetime, timedelta


class StockDataManager:
    def __init__(self, db_config: Dict[str, Any] = None):
        """
        初始化数据管理器

        Args:
            db_config: 数据库配置字典，可选参数包括:
                - db_path: 数据库文件路径 (默认: 'data/stock_data.db')
                - backup_path: 备份路径 (默认: 'data/backups/')
                - auto_backup: 是否自动备份 (默认: True)
        """
        # 数据库配置
        self.db_config = DatabaseConfig(db_config)
        self.db = StockDatabase(self.db_config)

        # 初始化数据源
        self.data_sources = {
            'akshare': AkShareDataSource(),
            'yahoo': YahooDataSource(),
            'baostock': BaostockDataSource()
        }

        self.logger = logging.getLogger(__name__)
        self._print_welcome_message()

    def _print_welcome_message(self):
        """打印欢迎信息和数据库状态"""
        print("=" * 50)
        print("股票数据管理系统")
        print("=" * 50)

        # 显示数据库信息
        db_info = self.db.get_database_info()
        print(f"数据库文件: {self.db_config.get_db_path()}")
        print(f"数据库大小: {db_info.get('db_size_mb', 0)} MB")
        print(f"数据表数量: {len(db_info.get('tables', []))}")

        for table, count in db_info.get('table_counts', {}).items():
            print(f"  - {table}: {count} 条记录")

        print("=" * 50)

    def show_database_status(self):
        """显示数据库状态"""
        db_info = self.db.get_database_info()

        print("\n=== 数据库状态 ===")
        print(f"路径: {self.db_config.get_db_path()}")
        print(f"大小: {db_info.get('db_size_mb', 0)} MB")
        print(f"表数量: {len(db_info.get('tables', []))}")

        for table, count in db_info.get('table_counts', {}).items():
            print(f"  {table}: {count} 条记录")

    def initialize_database(self, full_init: bool = False):
        """
        初始化数据库

        Args:
            full_init: 是否完全初始化（会清空现有数据）
        """
        if full_init:
            self.logger.warning("完全初始化将清空所有数据！")
            confirm = input("确认要清空所有数据吗？(y/N): ")
            if confirm.lower() != 'y':
                print("操作已取消")
                return

        self.logger.info("开始初始化数据库...")

        # 备份现有数据
        if os.path.exists(self.db_config.get_db_path()) and full_init:
            self.db.backup_database()

        # 重新初始化数据库
        self.db = StockDatabase(self.db_config)
        self.logger.info("数据库初始化完成")

    def run_maintenance(self):
        """运行数据库维护任务"""
        self.logger.info("开始数据库维护...")

        # 备份数据库
        if self.db_config.should_auto_backup():
            self.db.backup_database()

        # 优化数据库
        self.db.optimize_database()

        # 清理旧数据
        self.db.cleanup_old_data()

        self.logger.info("数据库维护完成")

    def update_stock_list(self, source: str = 'akshare'):
        """更新股票列表"""
        if source not in self.data_sources:
            self.logger.error(f"不支持的数据源: {source}")
            return

        data_source = self.data_sources[source]
        stock_list = data_source.get_stock_list()

        if not stock_list.empty:
            self.db.save_stock_basic(stock_list, source)
            self.logger.info(f"从 {source} 更新了 {len(stock_list)} 只股票信息")

    def update_daily_data(self, symbols: List[str] = None, source: str = 'akshare'):
        """更新日线数据"""
        if source not in self.data_sources:
            self.logger.error(f"不支持的数据源: {source}")
            return

        data_source = self.data_sources[source]

        # 如果没有指定股票，获取所有股票
        if symbols is None:
            conn = self.db._get_connection()
            symbols_df = pd.read_sql("SELECT symbol FROM stock_basic", conn)
            symbols = symbols_df['symbol'].tolist()
            conn.close()

        success_count = 0
        for i, symbol in enumerate(symbols):
            try:
                # 获取最后更新日期
                last_date = self.db.get_last_trade_date(symbol)
                start_date = (pd.to_datetime(last_date) + pd.Timedelta(days=1)).strftime('%Y-%m-%d')
                end_date = pd.Timestamp.now().strftime('%Y-%m-%d')

                if pd.to_datetime(start_date) > pd.to_datetime(end_date):
                    continue  # 数据已是最新

                df = data_source.get_daily_data(symbol, start_date, end_date)
                if not df.empty:
                    self.db.save_daily_quotes(df, source)
                    success_count += 1

                # 控制请求频率
                if (i + 1) % 10 == 0:
                    self.logger.info(f"已处理 {i+1}/{len(symbols)} 只股票")

            except Exception as e:
                self.logger.error(f"更新 {symbol} 数据失败: {e}")
                continue

        self.logger.info(f"日线数据更新完成，成功更新 {success_count}/{len(symbols)} 只股票")

    def get_data(self, symbol: str, start_date: str = None, end_date: str = None) -> pd.DataFrame:
        """获取股票数据"""
        return self.db.get_daily_data(symbol, start_date, end_date)






    def update_baostock_indicators(self, symbols: List[str] = None):
        """更新Baostock特有指标数据"""
        baostock_source = self.data_sources['baostock']

        if symbols is None:
            # 获取A股所有股票
            symbols_df = baostock_source.get_stock_list()
            symbols = symbols_df['symbol'].tolist()

        success_count = 0
        for i, symbol in enumerate(symbols):
            try:
                # 获取最近一年的数据
                end_date = datetime.now().strftime('%Y-%m-%d')
                start_date = (datetime.now() - timedelta(days=365)).strftime('%Y-%m-%d')

                df = baostock_source.get_daily_data(symbol, start_date, end_date)
                if not df.empty:
                    # 提取指标数据
                    indicator_cols = ['symbol', 'trade_date', 'pe_ratio', 'pb_ratio',
                                      'turnover_rate', 'change_rate']
                    indicator_df = df[['symbol', 'trade_date']].copy()

                    # 添加其他指标（这里可以根据需要扩展）
                    for col in ['pe_ratio', 'pb_ratio']:
                        if col in df.columns:
                            indicator_df[col] = df[col]

                    if not indicator_df.empty:
                        self.db.save_baostock_indicators(indicator_df)
                        success_count += 1

                # 控制请求频率
                if (i + 1) % 5 == 0:
                    time.sleep(1)  # Baostock建议每秒最多5次请求

                if (i + 1) % 100 == 0:
                    self.logger.info(f"已处理 {i+1}/{len(symbols)} 只股票的指标数据")

            except Exception as e:
                self.logger.error(f"更新 {symbol} 指标数据失败: {e}")
                continue

        self.logger.info(f"Baostock指标数据更新完成，成功更新 {success_count} 只股票")

    def update_dividend_data(self, symbols: List[str] = None, year: str = None):
        """更新分红数据"""
        baostock_source = self.data_sources['baostock']

        if symbols is None:
            # 获取大盘蓝筹股
            symbols = ['000001', '600036', '601318', '600519', '000858']

        success_count = 0
        for symbol in symbols:
            try:
                df = baostock_source.get_dividend_data(symbol, year)
                if not df.empty:
                    self.db.save_dividend_data(df, 'baostock')
                    success_count += 1
                    self.logger.info(f"更新 {symbol} 分红数据成功，共 {len(df)} 条记录")
            except Exception as e:
                self.logger.error(f"更新 {symbol} 分红数据失败: {e}")
                continue

        self.logger.info(f"分红数据更新完成，成功更新 {success_count} 只股票")